# 导入库
import pandas as pd
import statsmodels.formula.api as smf
from sqlalchemy import create_engine
import pymysql
from matplotlib import pyplot as plt
import seaborn as sns
import numpy as np

db_config = {
    'host': 'localhost',
    'user': 'root',
    'password': 'sjk1234',
    'database': 'tushare',
    'port': 3306,
    'charset': 'utf8'
}

engine = create_engine(f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}?charset={db_config['charset']}")
conn = pymysql.connect(**db_config)
chunk_size = 10000

df = pd.read_sql_query("SELECT * FROM date_1_d WHERE d.the_date BETWEEN '2023-01-01' AND '2023-12-31' and d.ts_code = '000001.SZ'",
                      conn,
                      chunksize=chunk_size
                      )
df1 = pd.concat(df, ignore_index=True)
print(df1.head())


df1['2d_closes'] = round((df1['closes'] - df1['closes'].shift(1)) / df1['closes'].shift(1), 2)
print(df1.head())

# 处理缺失数据
df1 = df1.dropna(subset=['2d_closes']).reset_index(drop=True)
print(df1.head())

ex = ['id', 'ts_code', 'trade_date', 'the_date', 'opens', 'high', 'low', 'closes', 'pre_closes', 'changes', 'pct_chg', 'i_closes']
number = df1.select_dtypes(include=['number']).columns.tolist()
newList = [i for i in number if i not in ex]

#选择需要进行主成分分析的自变量
features = df1[['vol', 'buy_sm_vol', 'sell_sm_vol', 'buy_md_vol', 'sell_md_vol', 'buy_lg_vol', 'sell_lg_vol', 'buy_elg_vol', 'sell_elg_vol', 'net_mf_vol']]
#计算特征值和特征向量
eigenvalues, eigenvectors = np.linalg.eig(X.corr())
print('累计贡献率：', round(eigenvalues[:5].sum() / eigenvalues.sum(), 4)* 100, '%')

n_components = 5
top_eigenvectors = eigenvectors[:, :n_components]

#计算主成分
principal_components = np.dot(x, top_eigenvectors)

#将主成分添加到数据框中
date_pca = pd.concat([df1, pd.DataFrame(principal_components, columns=[f'PC{i+1}' for i in range(n_components)])], axis=1)

#添加常数列前确保数据类型正确
X_pac = date_pca[[f'PC{i+1}' for i in range(n_components)]].copy()
X_pac = sm.add_constant(X_pac)

#确保Y的索引与X_pac一致
Y_pac = date_pca['2d_closes'].copy()

# 构建回归模型
print(X_pca)
print(y)
model = sm.OLS(y, X_pca)

# 拟合模型
results = model.fit()

# 输出结果
print('回归模型结果：')
print(results.summary())

# 选取PC3、PC4、PC5作为新的自变量
X_pca_selected = data_pca[['PC3', 'PC4', 'PC5']]
X_pca_selected.columns = ['PC3', 'PC4', 'PC5']

# 添加常数项
X_pca_selected = sm.add_constant(X_pca_selected)

# 因变量
y = data_pca['zd_pct_chg'].copy()

# 构建回归模型
model = sm.OLS(y, X_pca_selected)

# 拟合模型
results = model.fit()

# 输出结果
print('回归模型结果：')
print(results.summary())
X_pca_selected = data_pca[['PC3', 'PC4', 'PC5']]
X_pca_selected.columns = ['PC3', 'PC4', 'PC5']

# 绘制散点图
fig, axes = plt.subplots(1, 4, figsize=(15, 5))
for i, col in enumerate(X_pca_selected.columns):
    axes[i].scatter(X_pca_selected[col], y, s=50, alpha=0.7)
    axes[i].set_xlabel(col)
    axes[i].set_ylabel('zd_pct_chg')
    axes[i].set_title(f'{col}')
plt.tight_layout()
plt.show()

# 计算相关系数
for k in range(5):
    string_y = f'CP{k+1} = '
    i = eigenvectors[k]
    for j in range(len(i)):
        if i[j] > 0:
            string_y = string_y + f' {round(i[j], 2)} *X_{j+1}'
        else:
            string_y = string_y + f' {round(abs(i[j]), 2)} *X_{j+1}'
    if k !=2 and k !=4:
        print(string_y)
